/*
 * Copyright 2024 ywcai
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

`include "defines.v"
`timescale 1ns/1ps

module ex(
    input   wire                clk,
    input   wire                rst_n,

	input	wire				stall_ex_req_i,
	input	wire				stall_mem_req_i,
	input	wire[`MemAddrBus]	pc_i,
	input	wire[`InstDataBus]	inst_i,
	input	wire				i_inst_auipc_i,
	input	wire    			mux_alu_src1_i,
	input	wire				mux_alu_src2_i,
	input	wire[3:0]			mux_alu_ctrl_i,
    input   wire[3:0]           mux_mul_ctrl_i,
	input	wire				mem_req_i,
	input	wire				mem_we_i,
	input	wire				mux_branch_i,
	input	wire[2:0]			mux_result_src_i,
	input	wire				mux_jump_i,
	input	wire				mux_pctarget_src_i,
	input	wire				mux_loadtype_i,
	input	wire				mux_storetype_i,
	input	wire[`RegAddrBus]	rs1_addr_i,
	input	wire[`RegDataBus]	rs1_data_i,
	input	wire[`RegAddrBus]	rs2_addr_i,
	input	wire[`RegDataBus]	rs2_data_i,
	input	wire				regwrite_en_i,
	input	wire[`RegAddrBus]	regwrite_addr_i,
	input	wire[`RegDataBus]	rs1_imm_value_i,
	input	wire[`RegDataBus]	rs2_imm_value_i,
	input	wire[1:0]			mux_forward_rs1_i,
	input	wire[1:0]			mux_forward_rs2_i,
	input	wire				w_suffix_i,
	input	wire[`RegDataBus]	wb_regwrite_data_i,
	input	wire[`RegDataBus]	mem_forward_result_i,
	input	wire[`MemAddrBus]	mem_forward_pctarget_i,
	input	wire[`TrapBus]		trap_code_i,
    input   wire[`RegAddrBus]   mem_regwrite_addr_i,
    input   wire                mem_inst_mulh_i,
    input   wire                mem_inst_div_i,
    input   wire[`RegDataBus]   mul_low_result_i,
    input   wire[`RegDataBus]   div_rem_result_i,
	input	wire				req_grant_i,

	output	wire[`MemAddrBus]	pc_o,
	output	wire[`InstDataBus]	inst_o,
	output	wire				i_inst_auipc_o,
	output	wire				regwrite_en_o,
	output  wire[`RegAddrBus]	regwrite_addr_o,
	output	wire				mem_req_o,
	output	wire				mem_we_o,
	output	wire[`MemAddrBus]	mem_addr_o,
	output	wire[7:0]			mem_strb_o,
	output	wire[`RegDataBus]	mem_wr_data_o,
	output	wire[`RegDataBus]	rs1_data_o,
	output	wire[`RegDataBus]	rs2_data_o,
	output	wire[`RegDataBus]	rs1_imm_value_o,
    output  wire[`RegDataBus]   ex_result_o,
	output	wire[2:0]			mux_result_src_o,
	output	wire				mux_pcsrc_o,
	output	wire[`MemAddrBus]	pc_target_o,
	output	wire[`TrapBus]		trap_code_o,
    output  wire                mem_inst_mulh_o,
    output  wire                mem_inst_div_o,
    output  wire[`RegDataBus]   mul_low_result_o,
    output  wire[`RegDataBus]   div_rem_result_o,
    output  wire                div_stall_req_o,
	output	wire				mem_req_stall_o
	);

	reg		req_issue, req_accepted;

	wire[6:0] opcode		= inst_i[6:0];
	wire[2:0] funct3		= inst_i[14:12];
	wire[6:0] funct7		= inst_i[31:25];

	assign pc_o				= pc_i;
	assign inst_o			= inst_i;
	assign i_inst_auipc_o	= i_inst_auipc_i;
	assign regwrite_en_o	= regwrite_en_i;
	assign regwrite_addr_o	= regwrite_addr_i;
	assign mem_req_o		= req_issue;
	assign mem_we_o			= mem_we_i;
	assign mux_result_src_o	= mux_result_src_i;
	assign rs1_imm_value_o	= rs1_imm_value_i;

	wire[`RegDataBus] reg1_src = ({`XLEN{mux_forward_rs1_i == 2'b00}} & rs1_data_i)
		| ({`XLEN{mux_forward_rs1_i == 2'b01}} & wb_regwrite_data_i)
		| ({`XLEN{mux_forward_rs1_i == 2'b10}} & mem_forward_result_i)
		| ({`XLEN{mux_forward_rs1_i == 2'b11}} & mem_forward_pctarget_i);
	wire[`RegDataBus] operand_1 = mux_alu_src1_i ? rs1_imm_value_i : reg1_src;
	assign rs1_data_o = reg1_src;

	wire[`RegDataBus] reg2_src = ({`XLEN{mux_forward_rs2_i == 2'b00}} & rs2_data_i)
		| ({`XLEN{mux_forward_rs2_i == 2'b01}} & wb_regwrite_data_i)
		| ({`XLEN{mux_forward_rs2_i == 2'b10}} & mem_forward_result_i)
		| ({`XLEN{mux_forward_rs2_i == 2'b11}} & mem_forward_pctarget_i);
	wire[`RegDataBus] operand_2 = mux_alu_src2_i ? rs2_imm_value_i : reg2_src;
	assign rs2_data_o = reg2_src;

	wire[9:0] alu_op;
	assign alu_op[0] = (mux_alu_ctrl_i == 4'b0000);
	assign alu_op[1] = (mux_alu_ctrl_i == 4'b0001);
	assign alu_op[2] = (mux_alu_ctrl_i == 4'b0010);
	assign alu_op[3] = (mux_alu_ctrl_i == 4'b0011);
	assign alu_op[4] = (mux_alu_ctrl_i == 4'b0100);
	assign alu_op[5] = (mux_alu_ctrl_i == 4'b0101);
	assign alu_op[6] = (mux_alu_ctrl_i == 4'b0110);
	assign alu_op[7] = (mux_alu_ctrl_i == 4'b0111);
	assign alu_op[8] = (mux_alu_ctrl_i == 4'b1000);
	assign alu_op[9] = (mux_alu_ctrl_i == 4'b1001);

	wire[`RegDataBus] ex_result;
	wire z_flag;
	wire c_flag;
	wire n_flag;
	wire alu_v_flag;

	alu i_alu(
		.alu_op_i(alu_op),
		.operand_1_i(operand_1),
		.operand_2_i(operand_2),
		.w_suffix_i(w_suffix_i),
		.alu_res_o(ex_result),
		.z_flag_o(z_flag),
		.c_flag_o(c_flag),
		.n_flag_o(n_flag),
		.v_flag_o(alu_v_flag)
	);

	wire i_inst_auipc	= (opcode == `INST_AUIPC);

	wire i_inst_mul		= (mux_mul_ctrl_i == 4'b0001);
	wire i_inst_mulh	= (mux_mul_ctrl_i == 4'b0010);
	wire i_inst_mulhu	= (mux_mul_ctrl_i == 4'b0011);
	wire i_inst_mulhsu	= (mux_mul_ctrl_i == 4'b0100);
	wire i_inst_div		= (mux_mul_ctrl_i == 4'b0101);
	wire i_inst_divu	= (mux_mul_ctrl_i == 4'b0110);
	wire i_inst_rem		= (mux_mul_ctrl_i == 4'b0111);
	wire i_inst_remu	= (mux_mul_ctrl_i == 4'b1000);
	wire i_inst_m_cat	= (opcode == `R_TYPE || opcode == `RW_TYPE) && funct7[0];

	assign mem_inst_mulh_o = i_inst_mulh | i_inst_mulhu | i_inst_mulhsu;
	wire[`DXLEN-1:0] mul_result;
	wire mul_fuse = mem_inst_mulh_i & i_inst_mul
		& (mem_regwrite_addr_i != regwrite_addr_i);

	mul i_mul(
		.mux_mul_ctrl_i(mux_mul_ctrl_i),
		.w_suffix_i(w_suffix_i),
		.mul_operand_1_i(operand_1),
		.mul_operand_2_i(operand_2),
		.mul_result_o(mul_result)
	);
	assign mul_low_result_o = mul_result[`XLEN-1:0];

	wire[`RegDataBus] div_res_quo;
	wire[`RegDataBus] div_res_rem;
	wire div_overflow;
	assign mem_inst_div_o = i_inst_div | i_inst_divu;
	wire i_inst_div_cat = i_inst_div | i_inst_divu
		| i_inst_rem | i_inst_remu;
	wire div_fuse = mem_inst_div_i & (i_inst_rem | i_inst_remu)
		& (mem_regwrite_addr_i != regwrite_addr_i);
	wire div_req = i_inst_div_cat & !div_fuse;

	div i_div(
		.clk(clk),
		.rst_n(rst_n),
		.div_req_i(div_req),
		.mux_mul_ctrl_i(mux_mul_ctrl_i),
		.div_operand_1_i(operand_1),
		.div_operand_2_i(operand_2),
		.div_quotient_o(div_res_quo),
		.div_remainder_o(div_res_rem),
		.v_flag_o(div_overflow),
		.div_busy_o(div_stall_req_o)
	);
	assign div_rem_result_o = div_res_rem;

	wire[`MemAddrBus] pc_target;
	make_pctarget i_make_pctarget(
		.pc_i(pc_i),
		.offset_i(rs2_imm_value_i),
		.pc_target_o(pc_target)
	);

	assign ex_result_o = ({`XLEN{!i_inst_m_cat && !i_inst_auipc}} & ex_result)
		| ({`XLEN{i_inst_auipc}} & {{24{pc_target[39]}},pc_target})
		| ({`XLEN{!mul_fuse && i_inst_mul && !w_suffix_i}} & mul_result[`DATA_MSB:0])
		| ({`XLEN{i_inst_mul && w_suffix_i}} & {{32{mul_result[31]}}, mul_result[31:0]})
		| ({`XLEN{i_inst_mulh || i_inst_mulhu || i_inst_mulhsu}} & mul_result[`DXLEN-1:`XLEN])
		| ({`XLEN{i_inst_div || i_inst_divu}} & div_res_quo)
		| ({`XLEN{!div_fuse && (i_inst_rem || i_inst_remu)}} & div_res_rem)
		| ({`XLEN{mul_fuse}} & mul_low_result_i)
		| ({`XLEN{div_fuse}} & div_rem_result_i);

	assign pc_target_o = mux_pctarget_src_i ? {ex_result[39:1],1'b0} : pc_target;

	wire unsigned_lt = ~c_flag;
	wire signed_lt = n_flag ^ alu_v_flag;
	wire beq_jump = mux_branch_i & (funct3 == `FUNCT3_000) & z_flag;
	wire bne_jump = mux_branch_i & (funct3 == `FUNCT3_001) & !z_flag;
	wire blt_jump = mux_branch_i & (funct3 == `FUNCT3_100) & (signed_lt);
	wire bge_jump = mux_branch_i & (funct3 == `FUNCT3_101) & (~signed_lt);
	wire bltu_jump = mux_branch_i & (funct3 == `FUNCT3_110) & unsigned_lt;
	wire bgeu_jump = mux_branch_i & (funct3 == `FUNCT3_111) & (~unsigned_lt);

	assign mux_pcsrc_o = mux_jump_i | beq_jump | bne_jump
		| blt_jump | bge_jump | bltu_jump | bgeu_jump;

	assign mem_addr_o = ex_result_o;

	wire i_inst_store = (opcode == `INST_STORE);
	wire i_inst_sb = (i_inst_store & (funct3 == `FUNCT3_000));
	wire i_inst_sh = (i_inst_store & (funct3 == `FUNCT3_001));
	wire i_inst_sw = (i_inst_store & (funct3 == `FUNCT3_010));
	wire i_inst_sd = (i_inst_store & (funct3 == `FUNCT3_011));
	wire i_inst_load = (opcode == `INST_LOAD);
	wire i_inst_lb	= (i_inst_load & (funct3 == `FUNCT3_000));
	wire i_inst_lbu	= (i_inst_load & (funct3 == `FUNCT3_100));
	wire i_inst_lh	= (i_inst_load & (funct3 == `FUNCT3_001));
	wire i_inst_lhu	= (i_inst_load & (funct3 == `FUNCT3_101));
	wire i_inst_lw	= (i_inst_load & (funct3 == `FUNCT3_010));
	wire i_inst_lwu	= (i_inst_load & (funct3 == `FUNCT3_110));
	wire i_inst_ld	= (i_inst_load & (funct3 == `FUNCT3_011));

	wire[2:0] byte_idx = mem_addr_o[2:0] & 3'b111;
	assign mem_wr_data_o = ({`XLEN{i_inst_sb}} & {8{reg2_src[7:0]}})
		| ({`XLEN{i_inst_sh}} & {4{reg2_src[15:0]}})
		| ({`XLEN{i_inst_sw}} & {2{reg2_src[31:0]}})
		| ({`XLEN{i_inst_sd}} & reg2_src);
	wire[7:0] byte_strb = ({8{byte_idx == 3'b000}} & 8'b00000001)
		| ({8{byte_idx == 3'b001}} & 8'b00000010)
		| ({8{byte_idx == 3'b010}} & 8'b00000100)
		| ({8{byte_idx == 3'b011}} & 8'b00001000)
		| ({8{byte_idx == 3'b100}} & 8'b00010000)
		| ({8{byte_idx == 3'b101}} & 8'b00100000)
		| ({8{byte_idx == 3'b110}} & 8'b01000000)
		| ({8{byte_idx == 3'b111}} & 8'b10000000);
	wire[1:0] half_idx = mem_addr_o[2:1] & 2'b11;
	wire[7:0] half_strb = ({8{half_idx == 2'b00}} & 8'b00000011)
		| ({8{half_idx == 2'b01}} & 8'b00001100)
		| ({8{half_idx == 2'b10}} & 8'b00110000)
		| ({8{half_idx == 2'b11}} & 8'b11000000);
	wire word_idx = mem_addr_o[2] & 1'b1;
	wire[7:0] word_strb = ({8{word_idx == 1'b0}} & 8'b00001111)
		| ({8{word_idx == 1'b1}} & 8'b11110000);
	assign mem_strb_o =
		({8{i_inst_sb | i_inst_lb | i_inst_lbu}} & byte_strb)
		| ({8{i_inst_sh | i_inst_lh | i_inst_lhu}} & half_strb)
		| ({8{i_inst_sw | i_inst_lw | i_inst_lwu}} & word_strb)
		| ({8{i_inst_sd | i_inst_ld}} & 8'b11111111);

	// temporary disabled for jalr-misalign1 test case
	wire inst_addr_misalign = `DISABLE/*mux_pcsrc_o && (pc_target_o[1:0] != 2'b00)*/;
	wire data_addr_misalign = ((i_inst_sh|i_inst_lh|i_inst_lhu)  && ex_result_o[0])
		|| ((i_inst_sw|i_inst_lw|i_inst_lwu) && (ex_result_o[1:0] != 2'b00))
		|| ((i_inst_sd|i_inst_ld) && (ex_result_o[2:0] != 3'b000));

	// ignore overflow alu_v_flag|div_overflow
	assign trap_code_o = {trap_code_i[15], data_addr_misalign
		, inst_addr_misalign, trap_code_i[12:0]};

	wire ex_readygo;
	wire ex_allowin;
	wire mem_allowin;

	assign ex_readygo = mem_req_o && req_grant_i;
	assign mem_allowin = !stall_mem_req_i;
	assign ex_allowin = ex_readygo && mem_allowin;
	assign mem_req_stall_o = mem_req_i && !ex_allowin;

	always @(*) begin
		if (rst_n == `RESET_ENABLE) begin
			req_issue = `DISABLE;
			req_accepted = `DISABLE;
		end else if (mem_req_i) begin
			// means pipeline stalled either because i-fetch or d-fetch
			if (stall_ex_req_i) begin
				if (!req_accepted) begin
					req_issue = `ENABLE;
				end else if (req_grant_i) begin
					req_issue = `DISABLE;
					req_accepted = `ENABLE;
				end
			end else if (!req_grant_i)begin
				req_issue = `ENABLE;
			end else begin
				req_accepted = `ENABLE;
			end
		end else begin
			req_issue = `DISABLE;
			req_accepted = `DISABLE;
		end
	end

endmodule
